import dns.resolver
import time
import threading
import queue
import argparse
from concurrent.futures import ThreadPoolExecutor
import statistics
import logging
import os
from datetime import datetime
from tqdm import tqdm

class DNSBenchmark:
    def __init__(self, dns_server, domains, threads=10, queries_per_thread=100):
        self.dns_server = dns_server
        self.domains = domains
        self.threads = threads
        self.queries_per_thread = queries_per_thread
        self.results = queue.Queue()
        self.resolver = dns.resolver.Resolver()
        self.resolver.nameservers = [dns_server]
        self.resolver.timeout = 2
        self.resolver.lifetime = 2
        
        # 进度跟踪
        self.total_queries = threads * queries_per_thread * len(domains)
        self.completed_queries = 0
        self.progress_lock = threading.Lock()
        self.progress_bar = tqdm(
            total=self.total_queries,
            desc="DNS查询进度",
            unit="次",
            ncols=100,
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
        )
        
        # 设置日志
        self.setup_logging()

    def setup_logging(self):
        # 创建 logs 目录
        os.makedirs('logs', exist_ok=True)
        
        # 生成日志文件名，包含时间戳
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = f'logs/dns_benchmark_{timestamp}.log'
        
        # 配置日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()  # 同时输出到控制台
            ]
        )
        self.logger = logging.getLogger(__name__)

    def update_progress(self):
        with self.progress_lock:
            self.completed_queries += 1
            self.progress_bar.update(1)
            if self.completed_queries % 100 == 0 or self.completed_queries == self.total_queries:
                progress = (self.completed_queries / self.total_queries) * 100
                self.logger.info(f"进度: {self.completed_queries}/{self.total_queries} ({progress:.2f}%)")

    def __del__(self):
        if hasattr(self, 'progress_bar'):
            self.progress_bar.close()

    def query_domain(self, domain):
        start_time = time.time()
        try:
            answers = self.resolver.resolve(domain, 'A')
            elapsed = (time.time() - start_time) * 1000  # 转换为毫秒
            result = {
                'domain': domain,
                'status': 'success',
                'time': elapsed,
                'answers': [str(rdata) for rdata in answers]
            }
        except Exception as e:
            elapsed = (time.time() - start_time) * 1000
            result = {
                'domain': domain,
                'status': 'failed',
                'time': elapsed,
                'error': str(e)
            }
        
        self.update_progress()
        return result

    def worker(self):
        for _ in range(self.queries_per_thread):
            for domain in self.domains:
                result = self.query_domain(domain)
                self.results.put(result)

    def run_benchmark(self):
        self.logger.info("\n开始 DNS 压测:")
        self.logger.info(f"DNS 服务器: {self.dns_server}")
        self.logger.info(f"并发线程数: {self.threads}")
        self.logger.info(f"每个线程查询次数: {self.queries_per_thread}")
        self.logger.info(f"测试域名: {', '.join(self.domains)}")
        self.logger.info(f"总查询次数: {self.total_queries}")
        self.logger.info("\n正在测试...\n")

        start_time = time.time()
        
        with ThreadPoolExecutor(max_workers=self.threads) as executor:
            futures = [executor.submit(self.worker) for _ in range(self.threads)]
            for future in futures:
                future.result()

        total_time = time.time() - start_time
        
        # 收集结果
        results_list = []
        success_times = []
        failed_count = 0
        
        while not self.results.empty():
            result = self.results.get()
            results_list.append(result)
            if result['status'] == 'success':
                success_times.append(result['time'])
            else:
                failed_count += 1

        # 计算统计数据
        total_queries = len(results_list)
        success_count = len(success_times)
        
        if success_times:
            avg_time = statistics.mean(success_times)
            min_time = min(success_times)
            max_time = max(success_times)
            p95_time = sorted(success_times)[int(len(success_times) * 0.95)]
        else:
            avg_time = min_time = max_time = p95_time = 0

        # 记录结果
        self.logger.info("\n测试结果:")
        self.logger.info(f"总查询次数: {total_queries}")
        self.logger.info(f"成功查询: {success_count}")
        self.logger.info(f"失败查询: {failed_count}")
        self.logger.info(f"总耗时: {total_time:.2f} 秒")
        self.logger.info(f"QPS: {total_queries/total_time:.2f}")
        self.logger.info("\n响应时间统计 (毫秒):")
        self.logger.info(f"平均: {avg_time:.2f}")
        self.logger.info(f"最小: {min_time:.2f}")
        self.logger.info(f"最大: {max_time:.2f}")
        self.logger.info(f"95分位: {p95_time:.2f}")

        # 记录错误样本
        if failed_count > 0:
            self.logger.info("\n错误示例:")
            error_samples = [r for r in results_list if r['status'] == 'failed'][:3]
            for sample in error_samples:
                self.logger.info(f"域名: {sample['domain']}, 错误: {sample['error']}")

def main():
    parser = argparse.ArgumentParser(description='DNS 服务器压测工具')
    parser.add_argument('--server', type=str, required=True, help='DNS 服务器 IP')
    parser.add_argument('--threads', type=int, required=True, help='并发线程数')
    parser.add_argument('--queries', type=int, required=True, help='每个线程的查询次数')
    parser.add_argument('--domains', type=str, required=True, help='要查询的域名，用逗号分隔')

    args = parser.parse_args()
    
    # 打印接收到的参数
    print("接收到的参数:")
    print(f"Server: {args.server}")
    print(f"Threads: {args.threads}")
    print(f"Queries: {args.queries}")
    print(f"Domains: {args.domains}")
    
    domains = [d.strip() for d in args.domains.split(',')]

    try:
        benchmark = DNSBenchmark(
            dns_server=args.server,
            domains=domains,
            threads=args.threads,
            queries_per_thread=args.queries
        )
        
        benchmark.run_benchmark()
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    finally:
        if 'benchmark' in locals():
            benchmark.progress_bar.close()

if __name__ == '__main__':
    main() 